{
    "cells": [
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "# Agentic RAG with Autogen using Azure AI Services\n",
                "\n",
                "This notebook demonstrates implementing Retrieval-Augmented Generation (RAG) using Autogen agents with enhanced evaluation capabilities."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 1,
            "metadata": {},
            "outputs": [],
            "source": [
                "import os\n",
                "import time\n",
                "import asyncio\n",
                "from typing import List, Dict\n",
                "\n",
                "from autogen_agentchat.agents import AssistantAgent\n",
                "from autogen_core import CancellationToken\n",
                "from autogen_agentchat.messages import TextMessage\n",
                "from azure.core.credentials import AzureKeyCredential\n",
                "from autogen_ext.models.azure import AzureAIChatCompletionClient\n",
                "\n",
                "from azure.search.documents import SearchClient\n",
                "from azure.search.documents.indexes import SearchIndexClient\n",
                "from azure.search.documents.indexes.models import SearchIndex, SimpleField, SearchFieldDataType, SearchableField\n"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## Create the Client \n",
                "\n",
                "First, we initialize the Azure AI Chat Completion Client. This client will be used to interact with the Azure OpenAI service to generate responses to user queries."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 2,
            "metadata": {},
            "outputs": [],
            "source": [
                "client = AzureAIChatCompletionClient(\n",
                "    model=\"gpt-4o-mini\",\n",
                "    endpoint=\"https://models.inference.ai.azure.com\",\n",
                "    credential=AzureKeyCredential(os.environ[\"GITHUB_TOKEN\"]),\n",
                "    model_info={\n",
                "        \"json_output\": True,\n",
                "        \"function_calling\": True,\n",
                "        \"vision\": True,\n",
                "        \"family\": \"unknown\",\n",
                "    },\n",
                ")"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## Vector Database Initialization\n",
                "\n",
                "We initialize Azure AI Search with persistent storage and add enhanced sample documents. Azure AI Search will be used to store and retrieve documents that provide context for generating accurate responses."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 3,
            "metadata": {},
            "outputs": [
                {
                    "data": {
                        "text/plain": [
                            "[<azure.search.documents._generated.models._models_py3.IndexingResult at 0x1c479d5b6b0>,\n",
                            " <azure.search.documents._generated.models._models_py3.IndexingResult at 0x1c479d5b710>,\n",
                            " <azure.search.documents._generated.models._models_py3.IndexingResult at 0x1c479d5b6e0>,\n",
                            " <azure.search.documents._generated.models._models_py3.IndexingResult at 0x1c479d5b740>,\n",
                            " <azure.search.documents._generated.models._models_py3.IndexingResult at 0x1c479d5b770>]"
                        ]
                    },
                    "execution_count": 3,
                    "metadata": {},
                    "output_type": "execute_result"
                }
            ],
            "source": [
                "# Initialize Azure AI Search with persistent storage\n",
                "search_service_endpoint = os.getenv(\"AZURE_SEARCH_SERVICE_ENDPOINT\")\n",
                "search_api_key = os.getenv(\"AZURE_SEARCH_API_KEY\")\n",
                "index_name = \"travel-documents\"\n",
                "\n",
                "search_client = SearchClient(\n",
                "    endpoint=search_service_endpoint,\n",
                "    index_name=index_name,\n",
                "    credential=AzureKeyCredential(search_api_key)\n",
                ")\n",
                "\n",
                "index_client = SearchIndexClient(\n",
                "    endpoint=search_service_endpoint,\n",
                "    credential=AzureKeyCredential(search_api_key)\n",
                ")\n",
                "\n",
                "# Define the index schema\n",
                "fields = [\n",
                "    SimpleField(name=\"id\", type=SearchFieldDataType.String, key=True),\n",
                "    SearchableField(name=\"content\", type=SearchFieldDataType.String)\n",
                "]\n",
                "\n",
                "index = SearchIndex(name=index_name, fields=fields)\n",
                "\n",
                "# Create the index\n",
                "index_client.create_index(index)\n",
                "\n",
                "# Enhanced sample documents\n",
                "documents = [\n",
                "    {\"id\": \"1\", \"content\": \"Contoso Travel offers luxury vacation packages to exotic destinations worldwide.\"},\n",
                "    {\"id\": \"2\", \"content\": \"Our premium travel services include personalized itinerary planning and 24/7 concierge support.\"},\n",
                "    {\"id\": \"3\", \"content\": \"Contoso's travel insurance covers medical emergencies, trip cancellations, and lost baggage.\"},\n",
                "    {\"id\": \"4\", \"content\": \"Popular destinations include the Maldives, Swiss Alps, and African safaris.\"},\n",
                "    {\"id\": \"5\", \"content\": \"Contoso Travel provides exclusive access to boutique hotels and private guided tours.\"}\n",
                "]\n",
                "\n",
                "# Add documents to the index\n",
                "search_client.upload_documents(documents)\n"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "def get_retrieval_context(query: str) -> str:\n",
                "    results = search_client.search(query)\n",
                "    context_strings = []\n",
                "    for result in results:\n",
                "        context_strings.append(f\"Document: {result['content']}\")\n",
                "    return \"\\n\\n\".join(context_strings) if context_strings else \"No results found\""
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## Agent Configuration\n",
                "\n",
                "We configure the retrieval and assistant agents. The retrieval agent is specialized in finding relevant information using semantic search, while the assistant generates detailed responses based on the retrieved information."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 5,
            "metadata": {},
            "outputs": [],
            "source": [
                "# Create agents with enhanced capabilities\n",
                "assistant = AssistantAgent(\n",
                "    name=\"assistant\",\n",
                "    model_client=client,\n",
                "    system_message=(\n",
                "        \"You are a helpful AI assistant that provides answers using ONLY the provided context. \"\n",
                "        \"Do NOT include any external information. Base your answer entirely on the context given below.\"\n",
                "    ),\n",
                ")"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## RAGEvaluator Class\n",
                "\n",
                "We define the `RAGEvaluator` class to evaluate the response based on various metrics like response length, source citations, response time, and context relevance."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 14,
            "metadata": {},
            "outputs": [],
            "source": [
                "class RAGEvaluator:\n",
                "    def __init__(self):\n",
                "        self.responses: List[Dict] = []\n",
                "\n",
                "    def evaluate_response(self, query: str, response: str, context: List[Dict]) -> Dict:\n",
                "        # Basic metrics: response length, citation count, and a simple relevance score.\n",
                "        start_time = time.time()\n",
                "        metrics = {\n",
                "            'response_length': len(response),\n",
                "            'source_citations': sum(1 for doc in context if doc[\"content\"] in response),\n",
                "            'evaluation_time': time.time() - start_time,\n",
                "            'context_relevance': self._calculate_relevance(query, context)\n",
                "        }\n",
                "        self.responses.append({\n",
                "            'query': query,\n",
                "            'response': response,\n",
                "            'metrics': metrics\n",
                "        })\n",
                "        return metrics\n",
                "\n",
                "    def _calculate_relevance(self, query: str, context: List[Dict]) -> float:\n",
                "        # Simple relevance score: fraction of the documents where the query appears.\n",
                "        return sum(1 for c in context if query.lower() in c[\"content\"].lower()) / len(context)"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## Query Processing with RAG\n",
                "\n",
                "We define the `ask_rag` function to send the query to the assistant, process the response, and evaluate it. This function handles the interaction with the assistant and uses the evaluator to measure the quality of the response."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 15,
            "metadata": {},
            "outputs": [],
            "source": [
                "async def ask_rag(query: str, evaluator: RAGEvaluator):\n",
                "    try:\n",
                "        retrieval_context = get_retrieval_context(query)\n",
                "        # Augment the query with the retrieval context.\n",
                "        augmented_query = (\n",
                "            f\"Retrieved Context:\\n{retrieval_context}\\n\\n\"\n",
                "            f\"User Query: {query}\\n\\n\"\n",
                "            \"Based ONLY on the above context, please provide the answer.\"\n",
                "        )\n",
                "\n",
                "        # Send the augmented query as a user message.\n",
                "        start_time = time.time()\n",
                "        response = await assistant.on_messages(\n",
                "            [TextMessage(content=augmented_query, source=\"user\")],\n",
                "            cancellation_token=CancellationToken(),\n",
                "        )\n",
                "        processing_time = time.time() - start_time\n",
                "\n",
                "        # Evaluate the response against our vector-store documents.\n",
                "        metrics = evaluator.evaluate_response(\n",
                "            query=query,\n",
                "            response=response.chat_message.content,\n",
                "            context=documents\n",
                "        )\n",
                "        return {\n",
                "            'response': response.chat_message.content,\n",
                "            'processing_time': processing_time,\n",
                "            'metrics': metrics,\n",
                "        }\n",
                "    except Exception as e:\n",
                "        print(f\"Error processing query: {e}\")\n",
                "        return None"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "# Example usage\n",
                "\n",
                "We initialize the evaluator and define the queries that we want to process and evaluate."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 16,
            "metadata": {},
            "outputs": [],
            "source": [
                "async def main():\n",
                "    evaluator = RAGEvaluator()\n",
                "    queries = [\n",
                "        \"Can you explain Contoso's travel insurance coverage?\", # Relevant document\n",
                "        \"What is Neural Network?\" # No relevant document\n",
                "    ]\n",
                "    for query in queries:\n",
                "        print(f\"\\nProcessing Query: {query}\")\n",
                "        result = await ask_rag(query, evaluator)\n",
                "        if result:\n",
                "            print(\"Response:\", result['response'])\n",
                "        print(\"\\n\" + \"=\"*60 + \"\\n\")"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## Run the Script\n",
                "\n",
                "We check if the script is running in an interactive environment or a standard script, and run the main function accordingly."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 17,
            "metadata": {},
            "outputs": [
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "\n",
                        "Processing Query: Can you explain Contoso's travel insurance coverage?\n",
                        "Response: Contoso's travel insurance covers medical emergencies, trip cancellations, and lost baggage.\n",
                        "\n",
                        "============================================================\n",
                        "\n",
                        "\n",
                        "Processing Query: What is Neural Network?\n",
                        "Response: No information is available on Neural Network based on the provided context.\n",
                        "\n",
                        "============================================================\n",
                        "\n"
                    ]
                }
            ],
            "source": [
                "if __name__ == \"__main__\":\n",
                "    if asyncio.get_event_loop().is_running():\n",
                "        await main()\n",
                "    else:\n",
                "        asyncio.run(main())"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": []
        }
    ],
    "metadata": {
        "kernelspec": {
            "display_name": "venv",
            "language": "python",
            "name": "python3"
        },
        "language_info": {
            "codemirror_mode": {
                "name": "ipython",
                "version": 3
            },
            "file_extension": ".py",
            "mimetype": "text/x-python",
            "name": "python",
            "nbconvert_exporter": "python",
            "pygments_lexer": "ipython3",
            "version": "3.12.4"
        }
    },
    "nbformat": 4,
    "nbformat_minor": 2
}
